import os
import json
import pickle
import numpy as np
from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection,
)
from typing import List, Dict, Any, Union, Tuple
from datetime import datetime
import shutil
from tqdm import tqdm
import glob
from PIL import Image

from models.clip_model import ClipModel

class KnowledgeBase:
    """知识库管理类，用于创建和查询基于图像的知识库"""
    
    def __init__(self, base_path: str = "./knowledge_base", clip_model: ClipModel = None, milvus_host: str = "localhost", milvus_port: str = "19530"):
        """
        初始化知识库
        
        Args:
            base_path: 知识库存储路径
            clip_model: CLIP模型实例
            milvus_host: Milvus服务器地址
            milvus_port: Milvus服务器端口
        """
        self.base_path = base_path
        self.clip_model = clip_model
        self.milvus_host = milvus_host
        self.milvus_port = milvus_port
        
        # 创建知识库目录
        os.makedirs(base_path, exist_ok=True)
        
        # 知识库元数据
        self.kb_info = {
            "product_types": [],
            "created_at": None,
            "updated_at": None,
            "total_images": 0,
            "features_dim": 0,
            "clip_model": None
        }
        
        # 初始化
        self.collection = None
        self.image_paths = []
        self.image_labels = []
        
        # 连接到Milvus
        try:
            connections.connect(host=milvus_host, port=milvus_port)
            print(f"已连接到Milvus服务器: {milvus_host}:{milvus_port}")
        except Exception as e:
            print(f"连接Milvus服务器失败: {str(e)}")
            print("请确保Milvus服务正在运行")
        
    def load_or_create(self, product_name: str) -> bool:
        """
        加载或创建产品知识库
        
        Args:
            product_name: 产品名称
            
        Returns:
            是否成功加载或创建
        """
        product_path = os.path.join(self.base_path, product_name)
        os.makedirs(product_path, exist_ok=True)
        
        # 检查是否存在元数据文件
        metadata_path = os.path.join(product_path, "metadata.json")
        paths_path = os.path.join(product_path, "image_paths.pkl")
        labels_path = os.path.join(product_path, "image_labels.pkl")
        
        collection_name = f"defect_kb_{product_name}"
        
        # 如果存在元数据文件，则加载
        if os.path.exists(metadata_path) and utility.has_collection(collection_name):
            try:
                # 加载元数据
                with open(metadata_path, 'r', encoding='utf-8') as f:
                    self.kb_info = json.load(f)
                
                # 加载图像路径和标签
                with open(paths_path, 'rb') as f:
                    self.image_paths = pickle.load(f)
                
                with open(labels_path, 'rb') as f:
                    self.image_labels = pickle.load(f)
                
                # 加载Milvus集合
                self.collection = Collection(collection_name)
                self.collection.load()
                
                print(f"已加载知识库: {product_name}")
                print(f"包含图片数: {self.kb_info['total_images']}")
                return True
            
            except Exception as e:
                print(f"加载知识库失败: {str(e)}")
                return False
        
        return True
    
    def build_from_folder(self, product_name: str, data_dir: str, batch_size: int = 32) -> bool:
        """
        从文件夹构建知识库
        
        Args:
            product_name: 产品名称
            data_dir: 数据文件夹路径（应包含良品和缺陷子文件夹）
            batch_size: 批处理大小
            
        Returns:
            是否成功构建
        """
        if self.clip_model is None:
            print("错误: 请先初始化CLIP模型")
            return False
        
        # 创建产品知识库目录
        product_path = os.path.join(self.base_path, product_name)
        os.makedirs(product_path, exist_ok=True)
        
        # 获取所有子文件夹（代表不同缺陷类型）
        subdirs = [d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))]
        
        if not subdirs:
            print(f"警告: {data_dir} 中没有找到子文件夹")
            return False
        
        all_features = []
        self.image_paths = []
        self.image_labels = []
        
        total_images = 0
        for subdir in subdirs:
            subdir_path = os.path.join(data_dir, subdir)
            image_files = []
            
            # 获取所有支持的图片文件
            for ext in ['png', 'jpg', 'jpeg']:
                image_files.extend(glob.glob(os.path.join(subdir_path, f"*.{ext}")))
            
            if not image_files:
                print(f"警告: {subdir_path} 中没有找到图片")
                continue
            
            print(f"处理 {subdir} 中的 {len(image_files)} 张图片...")
            
            # 批处理提取特征
            for i in range(0, len(image_files), batch_size):
                batch = image_files[i:i+batch_size]
                features, valid_paths = self.clip_model.extract_features(batch)
                
                if len(valid_paths) > 0:
                    all_features.append(features)
                    self.image_paths.extend(valid_paths)
                    self.image_labels.extend([subdir] * len(valid_paths))
                    total_images += len(valid_paths)
        
        if total_images == 0:
            print("错误: 没有成功处理任何图片")
            return False
        
        # 合并所有特征
        all_features = np.vstack(all_features)
        features_dim = all_features.shape[1]
        
        # 创建Milvus集合
        collection_name = f"defect_kb_{product_name}"
        
        # 如果集合已存在，则删除
        if utility.has_collection(collection_name):
            utility.drop_collection(collection_name)
        
        # 定义集合架构
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
            FieldSchema(name="label", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=features_dim)
        ]
        schema = CollectionSchema(fields, "产品表面缺陷知识库")
        self.collection = Collection(collection_name, schema)
        
        # 创建索引
        index_params = {
            "metric_type": "IP",  # 内积，对应余弦相似度
            "index_type": "IVF_FLAT",
            "params": {"nlist": 1024}
        }
        self.collection.create_index(field_name="embedding", index_params=index_params)
        
        # 插入数据
        data = [
            self.image_paths,
            self.image_labels,
            all_features.tolist()
        ]
        
        self.collection.insert(data)
        self.collection.load()
        
        # 更新知识库信息
        self.kb_info = {
            "product_name": product_name,
            "created_at": datetime.now().isoformat(),
            "updated_at": datetime.now().isoformat(),
            "total_images": total_images,
            "features_dim": features_dim,
            "clip_model": self.clip_model.model_name,
            "defect_types": list(set(self.image_labels))
        }
        
        # 保存索引和元数据
        self._save_knowledge_base(product_name)
        
        print(f"知识库构建完成: {product_name}")
        print(f"包含图片数: {total_images}")
        print(f"缺陷类型: {list(set(self.image_labels))}")
        
        return True
    
    def add_images(self, product_name: str, image_paths: List[str], labels: List[str]) -> bool:
        """
        向现有知识库中添加图片
        
        Args:
            product_name: 产品名称
            image_paths: 图片路径列表
            labels: 对应的标签列表
            
        Returns:
            是否成功添加
        """
        if len(image_paths) != len(labels):
            print("错误: 图片路径和标签数量不匹配")
            return False
        
        if self.clip_model is None:
            print("错误: 请先初始化CLIP模型")
            return False
        
        # 加载现有知识库
        if not self.load_or_create(product_name):
            print(f"警告: 无法加载知识库 {product_name}，将创建新知识库")
        
        # 提取特征
        features, valid_paths = self.clip_model.extract_features(image_paths)
        
        if len(valid_paths) == 0:
            print("错误: 没有成功处理任何图片")
            return False
        
        collection_name = f"defect_kb_{product_name}"
        
        # 如果集合不存在，创建新集合
        if not utility.has_collection(collection_name):
            features_dim = features.shape[1]
            
            # 定义集合架构
            fields = [
                FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
                FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
                FieldSchema(name="label", dtype=DataType.VARCHAR, max_length=64),
                FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=features_dim)
            ]
            schema = CollectionSchema(fields, "产品表面缺陷知识库")
            self.collection = Collection(collection_name, schema)
            
            # 创建索引
            index_params = {
                "metric_type": "IP",  # 内积，对应余弦相似度
                "index_type": "IVF_FLAT",
                "params": {"nlist": 1024}
            }
            self.collection.create_index(field_name="embedding", index_params=index_params)
            
            self.kb_info["features_dim"] = features_dim
            self.kb_info["clip_model"] = self.clip_model.model_name
        
        # 插入新数据
        valid_labels = [labels[image_paths.index(path)] for path in valid_paths]
        
        data = [
            valid_paths,
            valid_labels,
            features.tolist()
        ]
        
        self.collection.insert(data)
        
        # 更新图片路径和标签
        self.image_paths.extend(valid_paths)
        self.image_labels.extend(valid_labels)
        
        # 更新知识库信息
        self.kb_info["total_images"] = len(self.image_paths)
        self.kb_info["updated_at"] = datetime.now().isoformat()
        if "defect_types" in self.kb_info:
            defect_types = set(self.kb_info["defect_types"])
            defect_types.update(valid_labels)
            self.kb_info["defect_types"] = list(defect_types)
        else:
            self.kb_info["defect_types"] = list(set(valid_labels))
        
        # 保存索引和元数据
        self._save_knowledge_base(product_name)
        
        print(f"已向知识库 {product_name} 添加 {len(valid_paths)} 张图片")
        return True
    
    def search(self, query_image_path: str, top_k: int = 5, product_name: str = None) -> List[Dict]:
        """
        搜索最相似的图片
        
        Args:
            query_image_path: 查询图片路径
            top_k: 返回结果数量
            product_name: 产品名称，如果提供则限制在该产品的知识库中搜索
            
        Returns:
            最相似图片的信息列表
        """
        if self.collection is None:
            print("错误: 知识库未加载")
            return []
        
        if self.clip_model is None:
            print("错误: CLIP模型未初始化")
            return []
        
        # 如果指定了产品名称，确保加载该产品的知识库
        if product_name and self.kb_info.get("product_name") != product_name:
            if not self.load_or_create(product_name):
                print(f"错误: 无法加载知识库 {product_name}")
                return []
        
        # 提取查询图片特征
        query_features, _ = self.clip_model.extract_features([query_image_path])
        
        if query_features.shape[0] == 0:
            print(f"错误: 无法处理查询图片 {query_image_path}")
            return []
        
        # 搜索相似图片
        search_params = {
            "metric_type": "IP",
            "params": {"nprobe": 16}
        }
        
        # 增加搜索数量，以便后续处理
        search_limit = min(top_k * 3, 20)
        
        results = self.collection.search(
            data=[query_features.tolist()[0]],
            anns_field="embedding",
            param=search_params,
            limit=search_limit,
            output_fields=["image_path", "label"]
        )
        
        # 整理结果
        result_list = []
        for hits in results:
            for hit in hits:
                result_list.append({
                    "similarity": float(hit.score),
                    "path": hit.entity.get("image_path"),
                    "label": hit.entity.get("label")
                })
        
        # 优化结果：确保包含足够的good样本和缺陷样本
        good_samples = [r for r in result_list if r.get("label", "").lower() == "good"]
        defect_samples = [r for r in result_list if r.get("label", "").lower() != "good"]
        
        # 平衡结果集
        balanced_results = []
        
        # 检查最相似的结果是否为good
        is_likely_good = False
        if result_list and result_list[0].get("label", "").lower() == "good":
            is_likely_good = True
        
        # 如果最相似结果为good，增加good样本的权重
        if is_likely_good:
            # 优先添加good样本
            good_count = min(len(good_samples), int(top_k * 0.7))
            defect_count = min(len(defect_samples), top_k - good_count)
            
            balanced_results.extend(good_samples[:good_count])
            balanced_results.extend(defect_samples[:defect_count])
        else:
            # 确保结果中包含一定比例的good样本，以便比较
            good_count = min(len(good_samples), max(1, int(top_k * 0.3)))
            defect_count = min(len(defect_samples), top_k - good_count)
            
            balanced_results.extend(defect_samples[:defect_count])
            balanced_results.extend(good_samples[:good_count])
        
        # 如果结果不足，补充原始结果
        if len(balanced_results) < top_k and len(result_list) > len(balanced_results):
            # 找出尚未添加的结果
            used_paths = set(r.get("path") for r in balanced_results)
            remaining = [r for r in result_list if r.get("path") not in used_paths]
            balanced_results.extend(remaining[:top_k - len(balanced_results)])
        
        # 按相似度重新排序
        balanced_results.sort(key=lambda x: x.get("similarity", 0.0), reverse=True)
        
        return balanced_results[:top_k]
    
    def text_search(self, query_text: str, product_name: str, top_k: int = 5) -> List[Dict]:
        """
        使用文本查询搜索相关图片
        
        Args:
            query_text: 查询文本
            product_name: 产品名称
            top_k: 返回结果数量
            
        Returns:
            相关图片的信息列表
        """
        if self.collection is None:
            print("错误: 知识库未加载")
            return []
        
        if self.clip_model is None:
            print("错误: CLIP模型未初始化")
            return []
        
        # 确保知识库已加载
        if not self.load_or_create(product_name):
            print(f"错误: 无法加载知识库 {product_name}")
            return []
        
        try:
            # 使用CLIP模型将文本转换为特征向量
            query_features = self.clip_model.encode_text(query_text)
            
            if query_features is None or len(query_features) == 0 or query_features.shape[0] == 0:
                print(f"错误: 无法处理查询文本 {query_text}")
                return []
            
            # 搜索相关图片
            search_params = {
                "metric_type": "IP",
                "params": {"nprobe": 16}
            }
            
            # 增加搜索数量，以便后续处理
            search_limit = min(top_k * 3, 20)
            
            results = self.collection.search(
                data=[query_features.tolist()[0]],
                anns_field="embedding",
                param=search_params,
                limit=search_limit,
                output_fields=["image_path", "label"]
            )
            
            # 整理结果
            result_list = []
            for hits in results:
                for hit in hits:
                    result_list.append({
                        "similarity": float(hit.score),
                        "path": hit.entity.get("image_path"),
                        "label": hit.entity.get("label")
                    })
            
            # 优化结果：确保包含足够的good样本和缺陷样本
            good_samples = [r for r in result_list if r.get("label", "").lower() == "good"]
            defect_samples = [r for r in result_list if r.get("label", "").lower() != "good"]
            
            # 平衡结果集
            balanced_results = []
            
            # 如果查询与"good"或"正常"相关，增加good样本的权重
            if "good" in query_text.lower() or "正常" in query_text or "无缺陷" in query_text:
                # 优先添加good样本
                good_count = min(len(good_samples), int(top_k * 0.7))
                defect_count = min(len(defect_samples), top_k - good_count)
                
                balanced_results.extend(good_samples[:good_count])
                balanced_results.extend(defect_samples[:defect_count])
            else:
                # 确保结果中包含一定比例的good样本，以便比较
                good_count = min(len(good_samples), max(1, int(top_k * 0.3)))
                defect_count = min(len(defect_samples), top_k - good_count)
                
                balanced_results.extend(defect_samples[:defect_count])
                balanced_results.extend(good_samples[:good_count])
            
            # 如果结果不足，补充原始结果
            if len(balanced_results) < top_k and len(result_list) > len(balanced_results):
                # 找出尚未添加的结果
                used_paths = set(r.get("path") for r in balanced_results)
                remaining = [r for r in result_list if r.get("path") not in used_paths]
                balanced_results.extend(remaining[:top_k - len(balanced_results)])
            
            # 按相似度重新排序
            balanced_results.sort(key=lambda x: x.get("similarity", 0.0), reverse=True)
            
            return balanced_results[:top_k]
        except Exception as e:
            print(f"文本搜索失败: {str(e)}")
            # 如果文本搜索失败，返回一些随机的样本作为备用方案
            return self._get_random_samples(product_name, top_k)

    def _get_random_samples(self, product_name: str, count: int = 5) -> List[Dict]:
        """
        获取产品知识库中的随机样本
        
        Args:
            product_name: 产品名称
            count: 样本数量
            
        Returns:
            随机样本列表
        """
        try:
            # 确保知识库已加载
            if not self.load_or_create(product_name):
                return []
            
            if not self.image_paths or len(self.image_paths) == 0:
                return []
            
            # 从知识库中随机选择样本
            import random
            sample_indices = random.sample(range(len(self.image_paths)), min(count, len(self.image_paths)))
            
            result_list = []
            for idx in sample_indices:
                result_list.append({
                    "similarity": 0.0,  # 随机样本，相似度为0
                    "path": self.image_paths[idx],
                    "label": self.image_labels[idx]
                })
            
            return result_list
        except Exception as e:
            print(f"获取随机样本失败: {str(e)}")
            return []
    
    def get_defect_types(self, product_name: str) -> List[str]:
        """
        获取产品的缺陷类型列表
        
        Args:
            product_name: 产品名称
            
        Returns:
            缺陷类型列表
        """
        if not self.load_or_create(product_name):
            return []
        
        if "defect_types" in self.kb_info:
            return self.kb_info["defect_types"]
        else:
            return list(set(self.image_labels))
    
    def get_product_list(self) -> List[str]:
        """
        获取所有产品列表
        
        Returns:
            产品名称列表
        """
        products = []
        if os.path.exists(self.base_path):
            products = [d for d in os.listdir(self.base_path) 
                       if os.path.isdir(os.path.join(self.base_path, d))]
        
        return products
    
    def delete_knowledge_base(self, product_name: str) -> bool:
        """
        删除知识库
        
        Args:
            product_name: 产品名称
            
        Returns:
            是否成功删除
        """
        product_path = os.path.join(self.base_path, product_name)
        collection_name = f"defect_kb_{product_name}"
        
        try:
            # 删除Milvus集合
            if utility.has_collection(collection_name):
                utility.drop_collection(collection_name)
            
            # 删除本地文件
            if os.path.exists(product_path):
                shutil.rmtree(product_path)
                
            print(f"已删除知识库: {product_name}")
            
            # 重置当前加载的知识库（如果是被删除的知识库）
            if self.kb_info.get("product_name") == product_name:
                self.collection = None
                self.image_paths = []
                self.image_labels = []
                self.kb_info = {"product_types": []}
            
            return True
        except Exception as e:
            print(f"删除知识库失败: {str(e)}")
            return False
    
    def _save_knowledge_base(self, product_name: str) -> None:
        """
        保存知识库相关文件
        
        Args:
            product_name: 产品名称
        """
        product_path = os.path.join(self.base_path, product_name)
        os.makedirs(product_path, exist_ok=True)
        
        # 保存元数据
        metadata_path = os.path.join(product_path, "metadata.json")
        with open(metadata_path, 'w', encoding='utf-8') as f:
            json.dump(self.kb_info, f, ensure_ascii=False, indent=2)
        
        # 保存图像路径和标签
        paths_path = os.path.join(product_path, "image_paths.pkl")
        with open(paths_path, 'wb') as f:
            pickle.dump(self.image_paths, f)
        
        labels_path = os.path.join(product_path, "image_labels.pkl")
        with open(labels_path, 'wb') as f:
            pickle.dump(self.image_labels, f) 

    def get_defect_descriptions(self, product_name: str) -> Dict[str, str]:
        """
        获取产品各缺陷类型的详细描述
        
        Args:
            product_name: 产品名称
            
        Returns:
            缺陷类型及其描述的字典
        """
        # 确保知识库已加载
        if not self.load_or_create(product_name):
            return {}
        
        # 获取缺陷类型列表
        defect_types = self.get_defect_types(product_name)
        
        # 预定义的缺陷类型描述模板
        defect_descriptions = {
            # 通用缺陷类型
            "good": "无缺陷，产品表面完好无损，无任何异常特征。",
            "scratch": "划痕，表现为产品表面的线状损伤，通常呈现为细长的线条，可能有不同的深度和长度。",
            "crack": "裂纹，表现为产品材料的断裂或分离，通常比划痕更深更宽，可能呈现为不规则的线条或网状结构。",
            "spot": "斑点，表现为产品表面的局部颜色异常或材质变化，通常呈圆形或不规则形状。",
            "stain": "污渍，表现为产品表面的颜色异常区域，通常边缘模糊，可能是由液体或其他物质造成的。",
            "hole": "孔洞，表现为产品表面的穿透性损伤，通常呈现为圆形或不规则形状的凹陷或穿透。",
            "broken": "破损，表现为产品结构的严重损坏，可能包括碎裂、断裂或部分缺失。",
            "color": "颜色异常，表现为产品颜色与正常标准不符，可能是色差、褪色或变色。",
            "bent": "弯曲，表现为产品形状的变形，原本应该是直的或平的部分出现弯曲或扭曲。",
            "contamination": "污染，表现为产品表面附着异物或杂质，可能是灰尘、纤维或其他物质。",
            
            # 特定产品类型的缺陷
            "bent_wire": "线材弯曲，电线或金属线出现不正常的弯折或扭曲。",
            "broken_large": "大面积破损，产品有大块区域的结构性损坏或材料缺失。",
            "broken_small": "小面积破损，产品有小范围的结构性损坏或材料缺失。",
            "cable_swap": "电缆错位，电缆连接错误或位置异常。",
            "combined": "复合缺陷，产品同时存在多种类型的缺陷。",
            "cut": "切割损伤，产品表面有明显的切口或割痕。",
            "cut_inner_insulation": "内部绝缘层切割损伤，电缆或线材的内部绝缘层有切口或损伤。",
            "cut_outer_insulation": "外部绝缘层切割损伤，电缆或线材的外部绝缘层有切口或损伤。",
            "defective": "一般性缺陷，产品存在影响功能或外观的问题。",
            "fabric_border_defect": "织物边缘缺陷，织物产品边缘有磨损、松散或不规则。",
            "fabric_interior_defect": "织物内部缺陷，织物产品内部有织造错误、纹理异常或材质问题。",
            "faulty_imprint": "印刷缺陷，产品表面的印刷图案或文字有模糊、错位或缺失。",
            "flip": "翻转错误，产品安装方向或位置错误。",
            "glue": "胶水残留，产品表面有多余的胶水或粘合剂残留。",
            "liquid_contamination": "液体污染，产品表面有液体造成的污染或痕迹。",
            "metal_contamination": "金属污染，产品表面有金属碎屑或颗粒。",
            "misplaced": "位置错误，产品组件安装位置不正确。",
            "missing_cable": "缺少电缆，产品应有的电缆组件缺失。",
            "missing_part": "缺少部件，产品的某个组成部分缺失。",
            "missing_wire": "缺少线材，产品应有的电线或连接线缺失。",
            "oil": "油污，产品表面有油脂或润滑剂的污染。",
            "poke": "戳痕，产品表面有尖锐物体造成的点状凹陷或穿透。",
            "print": "印刷问题，产品表面的印刷有质量问题或错误。",
            "rough": "表面粗糙，产品表面不光滑，有明显的纹理或颗粒感。",
            "squeeze": "挤压变形，产品因受到挤压而导致的形状变形。",
            "thread": "螺纹问题，螺纹部分有损坏、变形或不完整。",
            "thread_side": "螺纹侧面问题，螺纹侧面有损坏或异常。",
            "thread_top": "螺纹顶部问题，螺纹顶部有损坏或异常。"
        }
        
        # 根据产品类型添加特定描述
        product_specific_descriptions = {
            "bottle": {
                "good": "瓶子表面完好无损，无划痕、裂纹或变形，颜色均匀，透明度正常。",
                "broken_large": "瓶子有大面积破损，可能影响瓶体结构完整性，出现明显的碎裂或缺失。",
                "broken_small": "瓶子有小面积破损，如瓶口或瓶底的小碎片缺失，不影响整体结构。",
                "contamination": "瓶子表面或内部有异物污染，可能是灰尘、纤维或其他物质附着。"
            },
            "cable": {
                "good": "电缆外观完好，绝缘层完整，无划痕、裂纹或变形，线材排列整齐。",
                "bent_wire": "电缆中的金属线材出现异常弯折或扭曲，影响电缆的正常使用。",
                "cut_inner_insulation": "电缆内部绝缘层有切口或损伤，可能导致短路或漏电。",
                "cut_outer_insulation": "电缆外部绝缘层有切口或损伤，露出内部导线或绝缘材料。"
            },
            "neu-det": {
                "good": "钢材表面完好无损，无任何缺陷，表面光滑平整，颜色均匀，无龟裂、夹杂物、斑块、点蚀或划痕等异常特征。正常的钢材表面应当具有一致的金属光泽，纹理均匀。",
                
                "crazing": "龟裂，表现为钢材表面的网状细小裂纹，呈现不规则的网格状或蜘蛛网状纹路。视觉特征：表面有密集的细小裂纹网络，通常覆盖一定区域，在光照下会反射出不同的光线。形成原因：通常由材料内部应力、热处理不当、快速冷却或长期使用中的疲劳导致。可能影响：龟裂会降低材料的机械强度，增加断裂风险，并可能成为腐蚀的起点，加速材料劣化。",
                
                "inclusion": "夹杂物，表现为钢材表面或内部的异物颗粒或非金属杂质。视觉特征：不规则形状的斑点、凸起或嵌入物，颜色通常与周围金属材质不同，可能呈现为深色点或亮点。形成原因：冶炼过程中未完全清除的杂质、氧化物、硫化物或其他非金属物质混入金属基体。可能影响：降低材料的纯度和强度，成为应力集中点，在高应力环境下可能导致裂纹扩展，影响产品的使用寿命和安全性。",
                
                "patches": "斑块，表现为钢材表面的不规则色差区域。视觉特征：局部区域颜色与周围金属表面不同，呈现出不规则的斑驳状，可能是较深或较浅的色调。形成原因：材料成分不均匀、局部氧化、腐蚀、表面处理不均或热处理温度不一致。可能影响：可能表明材料性能不均，影响美观性，严重情况下可能暗示潜在的结构问题或材料疲劳。",
                
                "pitted_surface": "点蚀表面，表现为钢材表面的小孔洞或凹坑。视觉特征：表面有大小不一的孔洞或凹坑，分布可能集中或分散，在光照下可见明显的阴影。形成原因：电化学腐蚀、酸性物质侵蚀、气泡形成后破裂或制造过程中的材料缺陷。可能影响：增加表面粗糙度，降低抗腐蚀性能，可能成为应力集中点，在高应力环境下导致裂纹形成，并影响产品的美观性和功能性。",
                
                "scratches": "划痕，表现为钢材表面的线状损伤。视觉特征：表面有清晰可见的线状痕迹，可能是直线或弧线形态，长度和深度各异，在光照下会反射出明显的线条。形成原因：硬物刮擦、加工过程中的工具痕迹、运输或安装过程中的机械损伤。可能影响：增加表面粗糙度，可能成为腐蚀起点，深度划痕可能导致应力集中，影响产品的美观性和在某些应用中的功能性。"
            },
            "capsule": {
                "good": "胶囊表面光滑完整，颜色均匀，无变形、破损或污染。",
                "crack": "胶囊表面有裂纹，可能导致内容物泄漏或胶囊完整性受损。",
                "faulty_imprint": "胶囊上的印刷文字或标记模糊、错位或缺失。",
                "scratch": "胶囊表面有划痕，表现为细线状损伤，但不影响胶囊密封性。"
            },
            "carpet": {
                "good": "地毯表面纹理均匀，颜色一致，无磨损、污渍或纤维脱落。",
                "color": "地毯颜色异常，可能是褪色、变色或染色不均。",
                "cut": "地毯表面有切割损伤，纤维断裂或缺失，形成明显的切口。",
                "hole": "地毯表面有孔洞，表现为局部纤维完全缺失，露出底层材料。"
            },
            "grid": {
                "good": "网格结构完整，线条均匀，交叉点连接牢固，无变形或断裂。",
                "bent": "网格出现弯曲变形，原本应该是直的线条出现弯折或扭曲。",
                "broken": "网格结构有断裂或破损，连接点可能分离或材料有缺失。",
                "glue": "网格表面或连接处有多余的胶水残留，影响外观或功能。"
            },
            "hazelnut": {
                "good": "榛子表面光滑，颜色自然，形状规则，无破损或霉变。",
                "crack": "榛子壳有裂纹，可能导致内核暴露或受污染。",
                "hole": "榛子表面有孔洞，可能是由虫蛀或机械损伤造成。",
                "print": "榛子表面有非自然的印记或痕迹，可能是加工过程中产生的。"
            },
            "leather": {
                "good": "皮革表面质地均匀，颜色一致，无划痕、裂纹或污渍。",
                "color": "皮革颜色异常，可能是褪色、变色或染色不均。",
                "cut": "皮革表面有切割损伤，形成明显的切口或割痕。",
                "fold": "皮革表面有折痕或皱褶，影响外观平整度。",
                "poke": "皮革表面有戳痕，表现为点状凹陷或穿透。"
            },
            "metal_nut": {
                "good": "金属螺母表面光滑，螺纹完整，无变形、腐蚀或毛刺。",
                "bent": "金属螺母出现弯曲变形，影响与螺栓的配合。",
                "color": "金属螺母表面颜色异常，可能是氧化、腐蚀或表面处理不良。",
                "flip": "金属螺母安装方向错误，螺纹朝向不正确。",
                "scratch": "金属螺母表面有划痕，表现为线状损伤。"
            },
            "pill": {
                "good": "药片表面光滑完整，颜色均匀，形状规则，无缺口或污染。",
                "color": "药片颜色异常，可能是变色、褪色或颜色不均。",
                "crack": "药片表面有裂纹，但未完全断裂。",
                "faulty_imprint": "药片上的印刷标记模糊、错位或缺失。",
                "pill_type": "药片类型错误，与标准规格不符。"
            },
            "screw": {
                "good": "螺丝表面光滑，螺纹完整，头部无损，无变形或腐蚀。",
                "manipulated_front": "螺丝前部被人为改动，可能影响螺纹功能。",
                "scratch_head": "螺丝头部有划痕，表现为线状损伤。",
                "scratch_neck": "螺丝颈部有划痕，可能影响与配件的配合。",
                "thread_damaged": "螺丝螺纹有损坏，可能是磨损、变形或断裂。"
            },
            "tile": {
                "good": "瓷砖表面光滑平整，颜色均匀，边缘规则，无缺口或裂纹。",
                "crack": "瓷砖表面有裂纹，可能贯穿整个瓷砖或仅在表面。",
                "glue_strip": "瓷砖表面有胶水条痕，可能是安装过程中残留的。",
                "gray_stroke": "瓷砖表面有灰色条纹或痕迹，影响美观。",
                "oil": "瓷砖表面有油污，表现为光亮的斑块或痕迹。"
            },
            "toothbrush": {
                "good": "牙刷刷毛排列整齐，手柄完整，颜色正常，无变形或损坏。",
                "defective": "牙刷有一般性缺陷，可能是刷毛不齐、变形或手柄有问题。",
                "color": "牙刷颜色异常，可能是褪色或变色。"
            },
            "transistor": {
                "good": "晶体管各引脚完整，封装无损，标记清晰，无变形或腐蚀。",
                "bent_lead": "晶体管引脚弯曲，不符合标准安装要求。",
                "cut_lead": "晶体管引脚被切断或损坏，长度不足或有断口。",
                "damaged_case": "晶体管封装受损，可能有裂纹、缺口或变形。",
                "misplaced": "晶体管安装位置错误，引脚排列或方向不正确。"
            },
            "wood": {
                "good": "木材表面光滑，纹理自然，无节疤、裂纹或虫蛀痕迹。",
                "color": "木材颜色异常，可能是变色、褪色或染色不均。",
                "hole": "木材表面有孔洞，可能是由虫蛀或机械损伤造成。",
                "liquid": "木材表面有液体痕迹，可能是水渍、油渍或其他液体造成的。",
                "scratch": "木材表面有划痕，表现为线状损伤。"
            },
            "zipper": {
                "good": "拉链齿排列整齐，拉头完整，运行顺畅，无缺齿或变形。",
                "broken_teeth": "拉链有断齿或缺齿，影响拉链的正常使用。",
                "fabric_border_defect": "拉链织物边缘有缺陷，可能是磨损、松散或不规则。",
                "rough": "拉链表面粗糙，运行不顺畅，可能有阻力或卡顿。",
                "split_teeth": "拉链齿分离或错位，无法正常咬合。"
            }
        }
        
        # 创建结果字典
        result = {}
        
        # 为每个缺陷类型添加描述
        for defect_type in defect_types:
            # 首先检查产品特定描述
            if product_name in product_specific_descriptions and defect_type in product_specific_descriptions[product_name]:
                result[defect_type] = product_specific_descriptions[product_name][defect_type]
            # 然后检查通用描述
            elif defect_type.lower() in defect_descriptions:
                result[defect_type] = defect_descriptions[defect_type.lower()]
            # 如果没有预定义描述，提供一个通用描述
            else:
                result[defect_type] = f"{defect_type}缺陷，表现为产品表面或结构的异常，需要检查和处理。"
        
        return result

    def get_defect_features_description(self, product_name: str) -> str:
        """
        获取产品缺陷特征的详细描述文本，用于提示词
        
        Args:
            product_name: 产品名称
            
        Returns:
            缺陷特征描述文本
        """
        # 获取缺陷类型及描述
        defect_descriptions = self.get_defect_descriptions(product_name)
        
        if not defect_descriptions:
            return ""
        
        # 构建描述文本
        description_text = f"{product_name}产品可能的缺陷类型及其特征:\n\n"
        
        for defect_type, description in defect_descriptions.items():
            description_text += f"- {defect_type}: {description}\n"
        
        return description_text 